#!/opt/local/bin/python

from linacDetailsReader import datareader
from rfvectors import cavity
import matplotlib.pyplot as plt
from numpy import array
from scipy import nan, pi
from globals import c_light as c

xlsData = datareader(fname = 'scLinacDetails.csv')
xlsData.calcT()
xlsData.calcE0()

numcavs = len(xlsData.data['ncells'])

Pfor, Pref, Pbeam = [], [], []
TAng, delF = [], []
for i in xrange(numcavs):
    if xlsData.data['beta0'][i] == 0.46:
        f = 352.21e6
        Ql = 2.37e5 / 1.5
        RoverQ = 500
        Eampfile = 'Ezvszspoke.dat'
    elif xlsData.data['beta0'][i] == 0.7:
        f = 704.42e6
        Ql = 8e5
        RoverQ = 300
        Eampfile = nan
    elif xlsData.data['beta0'][i] == 0.92:
        f = 704.42e6
        Ql = 7.5e5
        RoverQ = 477
        Eampfile = nan

    cavL = xlsData.data['ncells'][i] * xlsData.data['beta0'][i] * (c/f) / 2

    A = cavity(f=f, I_b=50e-3, V_cav=xlsData.E0T[i]*cavL, beamph=xlsData.data['phase'][i], RoverQ=RoverQ, Ql=Ql, 
            beta_0=xlsData.data['beta0'][i], beta=xlsData.data['beta'][i], Eampfile=Eampfile)
    #print i
    #A.drawvecs(plt.figure())
    #A.ax.set_xlim([-4e7, 4e7])
    #A.ax.set_ylim([-4e7, 4e7])
    #A.fig.savefig('test%02d.png' % i)
    Pfor.append(A.Pfor)
    Pref.append(A.Pref)
    Pbeam.append(A.Pbeam)
    TAng.append(A.TAng)
    delF.append(A.delF)
    print "Cavity #%d: Ql/Qlopt = %f :: RQscaling = %f" % (i, A.Ql/A.Ql0, A.RQscaling)
    #print "RQscaling = %f :: beta = %f" % (A.RQscaling, A.beta)

Pfor, Pref, Pbeam = array(Pfor), array(Pref), array(Pbeam)
TAng, delF = array(TAng), array(delF)

fig = plt.figure(1)
ax = fig.add_subplot(111)
ax.plot(Pfor/1e3, 'b', label='Pfor')
ax.plot(Pref/1e3, 'rx', label='Pref')
ax.plot(Pbeam/1e3, 'bx', label='Pbeam')
ax.set_xlim(right=numcavs+1)
ax.set_xlabel('Cavity')
ax.set_ylabel('Power (kW)')
ax.legend(loc=0)
ax.grid()
ax.axvline(x=35.5, color='g')
ax.axvline(x=99.5, color='g')

fig2 = plt.figure(2)
ax21 = fig2.add_subplot(121)
ax22 = fig2.add_subplot(122)
ax21.plot(TAng*180/pi, 'x')
ax21.set_ylabel('Tuning angle (degrees)')
ax22.plot(delF, 'x')
ax22.set_ylabel('Detuning frequency (Hz)')
for i in [ax21, ax22]:
    i.set_ylim(bottom=0)
    i.grid()
    i.axvline(x=35.5, color='g')
    i.axvline(x=99.5, color='g')

plt.show()

